
import torch
import numpy as np
import librosa
import noisereduce as nr
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')





class HuBERTTranscriptionConfig:


    SAMPLE_RATE = 16000
    ENABLE_NOISE_REDUCTION = True



    MODEL_NAME = "jonatasgrosman/wav2vec2-large-xlsr-53-french"


    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"✓ 使用设备: {self.device}")

config = HuBERTTranscriptionConfig()





class HuBERTModelLoader:

    def __init__(self):
        self._model = None
        self._processor = None

    @property
    def model(self):
        if self._model is None:
            print(f"📥 正在加载法语Wav2Vec2模型（首次会下载~378MB，轻量版）...")
            self._model = Wav2Vec2ForCTC.from_pretrained(
                config.MODEL_NAME
            ).to(config.device)
            self._model.eval()
            print("✓ 法语Wav2Vec2模型加载完成")
        return self._model

    @property
    def processor(self):
        if self._processor is None:
            print("📥 正在加载处理器...")
            self._processor = Wav2Vec2Processor.from_pretrained(
                config.MODEL_NAME
            )
            print("✓ 处理器加载完成")
        return self._processor

models = HuBERTModelLoader()





def load_and_preprocess_audio(audio_path, enable_noise_reduction=None):
    audio_path = Path(audio_path)
    if not audio_path.exists():
        raise FileNotFoundError(f"音频文件不存在: {audio_path}")

    print(f"📂 正在加载音频: {audio_path.name}")


    waveform, sr = librosa.load(
        str(audio_path),
        sr=config.SAMPLE_RATE,
        mono=True
    )

    duration = len(waveform) / config.SAMPLE_RATE
    print(f"✓ 音频加载成功: {duration:.2f} 秒")


    if enable_noise_reduction is None:
        enable_noise_reduction = config.ENABLE_NOISE_REDUCTION

    if enable_noise_reduction:
        print("🔧 正在降噪...")
        waveform = nr.reduce_noise(
            y=waveform,
            sr=config.SAMPLE_RATE,
            stationary=True,
            prop_decrease=1.0
        )
        print("✓ 降噪完成")


    max_val = np.abs(waveform).max()
    if max_val > 0:
        waveform = waveform / max_val * 0.95

    return waveform





def transcribe_with_hubert(waveform):
    print("🧠 正在使用 HuBERT/Wav2Vec2 提取特征并转录...")


    inputs = models.processor(
        waveform,
        sampling_rate=config.SAMPLE_RATE,
        return_tensors="pt",
        padding=True
    )


    input_values = inputs.input_values.to(config.device)


    with torch.no_grad():
        logits = models.model(input_values).logits


    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = models.processor.batch_decode(predicted_ids)[0]

    print(f"✓ HuBERT转录完成: {len(transcription)} 字符")
    return transcription





def hubert_transcription_pipeline(
    audio_path,
    enable_noise_reduction=None,
    return_details=False
):
    print("\n" + "="*60)
    print("🎯 HuBERT法语转录流程开始")
    print("="*60)

    try:

        waveform = load_and_preprocess_audio(audio_path, enable_noise_reduction)
        duration = len(waveform) / config.SAMPLE_RATE


        transcription = transcribe_with_hubert(waveform)

        print("\n" + "="*60)
        print("✅ HuBERT转录完成！")
        print("="*60)

        if return_details:
            return {
                'text': transcription,
                'waveform': waveform,
                'duration': duration,
                'device': str(config.device),
                'model': config.MODEL_NAME
            }
        else:
            return transcription

    except Exception as e:
        print(f"\n❌ HuBERT转录失败: {type(e).__name__}: {e}")
        raise





if __name__ == "__main__":
    print("""
    ╔══════════════════════════════════════════════════════════════╗
    ║          HuBERT法语转录系统 v2.0                              ║
    ║          使用 Wav2Vec2-CTC 法语模型                           ║
    ╚══════════════════════════════════════════════════════════════╝
    """)

    test_audio = "../音频仓库/测试mp3（French）.mp3"

    print(f"📌 测试文件: {test_audio}\n")

    try:
        result = hubert_transcription_pipeline(test_audio)

        print("\n" + "🎉 "*20)
        print(f"【转录结果】\n{result}")
        print("🎉 "*20)

    except FileNotFoundError:
        print(f"❌ 测试文件不存在: {test_audio}")
    except Exception as e:
        print(f"❌ 发生错误: {e}")
